'''
Created on Aug 19, 2009

@author: mkiyer
'''

import sys

class Gene(object):
    pass


def parse_kgtxinfo(fhd):
    kgtxinfo = {}
    for line in fhd:
        fields = line.strip().split('\t')
        kgtxinfo[fields[0]] = fields[1]
    return kgtxinfo

def parse_ucsc_genes(fhd):
    '''
    name varchar(255) NOT NULL default '',
    chrom varchar(255) NOT NULL default '',
    strand char(1) NOT NULL default '',
    txStart int(10) unsigned NOT NULL default '0',
    txEnd int(10) unsigned NOT NULL default '0',
    cdsStart int(10) unsigned NOT NULL default '0',
    cdsEnd int(10) unsigned NOT NULL default '0',
    exonCount int(10) unsigned NOT NULL default '0',
    exonStarts longblob NOT NULL,
    exonEnds longblob NOT NULL,
    proteinID varchar(40) NOT NULL default '',
    alignID varchar(255) NOT NULL default '',
    '''
    for line in fhd:
        fields = line.strip().split('\t')
        g = Gene()
        g.name = fields[0]
        g.chrom = fields[1]
        g.strand = fields[2]
        g.txstart = int(fields[3])
        g.txend = int(fields[4])
        g.cdsstart = int(fields[5])
        g.cdsend = int(fields[6])
        g.exon_count = int(fields[7])
        g.exon_starts = map(int, fields[8].split(',')[:-1])
        g.exon_ends = map(int, fields[9].split(',')[:-1])
        g.exons = zip(g.exon_starts, g.exon_ends)
        g.introns = zip(g.exon_ends, g.exon_starts[1:])
        yield g

def write_bed_line(fhd, chrom, start, end, name, score, strand):
    fhd.write('%s\t%d\t%d\t%s\t%s\t%s\n' %
              (chrom, start, end, name, score, strand)) 
    
if __name__ == "__main__":
    kgtxinfofile = sys.argv[1]
    kgfile = sys.argv[2]
    
    # get gene categories
    kgtxinfofhd = open(kgtxinfofile)
    kgtxinfo = parse_kgtxinfo(kgtxinfofhd)
    kgtxinfofhd.close()
        
    # get known genes
    # TODO: get from mysql server
    kgfhd = open(kgfile)
    genes = [g for g in parse_ucsc_genes(kgfhd)]
    
    # output files
    coding_exons_fhd = open('ucsc_coding_exons.bed', 'w')
    coding_introns_fhd = open('ucsc_coding_introns.bed', 'w')
    coding_tx_fhd = open('ucsc_coding_transcripts.bed', 'w')
    noncoding_fhd = open('ucsc_noncoding_transcripts.bed', 'w')
    
    for g in genes:
        if kgtxinfo[g.name] == 'coding':
            for i, e in enumerate(g.exons):
                write_bed_line(coding_exons_fhd, g.chrom, e[0], e[1], 
                               '%s_exon%d' % (g.name, i), '.', g.strand)
            # write the coding transcript
            write_bed_line(coding_tx_fhd, g.chrom, g.txstart, g.txend, 
                           g.name, '.', g.strand)
                 
            # 3' UTR
            if g.txstart != g.cdsstart:
                write_bed_line(coding_introns_fhd, g.chrom, g.txstart, g.cdsstart,
                               '%s_3UTR' % g.name, '.', g.strand)
            # introns
            for i, intron in enumerate(g.introns):
                if intron[0] != intron[1]:
                    write_bed_line(coding_introns_fhd, g.chrom, intron[0], intron[1], 
                                   '%s_intron%d' % (g.name, i), '.', g.strand)            
            # 5' UTR
            if g.cdsend != g.txend:
                write_bed_line(coding_introns_fhd, g.chrom, g.cdsend, g.txend,
                               '%s_5UTR' % g.name, '.', g.strand)            
        else:
            noncoding_fhd.write('%s\t%d\t%d\t%s\t.\t%s\n' % 
                                (g.chrom, g.txstart, g.txend, g.name, g.strand))
